// Author: 8891689
// https://github.com/8891689
#ifndef AES_MAES_H
#define AES_MAES_H

#include <stdint.h>
#include <stddef.h>
#include <immintrin.h> 

// 枚舉和宏定義
typedef enum {
    AES_KEY_128 = 4,
    AES_KEY_192 = 6,
    AES_KEY_256 = 8
} AES_KeySize;

typedef enum {
    AES_KEY_SCHEDULE_ENCRYPT,
    AES_KEY_SCHEDULE_DECRYPT
} AES_KeySchedule_Mode;

#define AES_BLOCK_SIZE 16
#define AES_Nb 4

#define AES_Nr_Rounds(Nk_enum_val) (6 + (Nk_enum_val))

// --- 核心接口 ---
int AES_KeyExpansion(const uint8_t *key, __m128i *round_keys, AES_KeySize Nk_enum, AES_KeySchedule_Mode mode);

// 單塊加解密 
void AES_Encrypt_Block(uint8_t *output_block, const uint8_t *input_block, const __m128i *round_keys, int num_rounds);
void AES_Decrypt_Block(uint8_t *output_block, const uint8_t *input_block, const __m128i *round_keys, int num_rounds);

// AVX2 優化的 ECB 模式 
__attribute__((target("aes,avx2")))
void AES_Encrypt_AVX2(uint8_t *blocks, size_t num_blocks, const __m128i *round_keys, int num_rounds);

__attribute__((target("aes,avx2")))
void AES_Decrypt_AVX2(uint8_t *blocks, size_t num_blocks, const __m128i *round_keys, int num_rounds);

void AES_PrintState(const uint8_t *state);

#endif
